"""
回测引擎
支持多交易所、多策略类型的独立回测系统
"""

import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any, Tuple
import logging
from pathlib import Path
import json
import matplotlib.pyplot as plt
import seaborn as sns

from ..core.unified_core import ConfigManager
from ..data.data_manager import DataManager
from ..utils.helpers import TimeUtils
from .mock_managers import MockOrderManager, MockAccountManager


class BacktestEngine:
    """回测引擎"""
    
    def __init__(self, config: Dict[str, Any]):
        """
        初始化回测引擎
        
        Args:
            config: 回测配置
        """
        self.config = config
        self.logger = logging.getLogger(__name__)
        
        # 回测参数
        self.start_date = config.get('start_date')
        self.end_date = config.get('end_date')
        self.initial_balance = config.get('initial_balance', 10000)
        self.commission_rate = config.get('commission_rate', 0.001)
        self.slippage_rate = config.get('slippage_rate', 0.0001)
        
        # 数据管理
        self.data_manager = DataManager(config.get('data', {}))
        
        # 模拟管理器
        self.order_manager = MockOrderManager()
        self.account_manager = MockAccountManager(self.initial_balance)
        
        # 回测结果
        self.results = {
            'trades': [],
            'positions': [],
            'balance_history': [],
            'metrics': {},
            'orders': []
        }
        
        # 当前状态
        self.current_time = None
        self.current_data = {}
        self.positions = {}
        self.strategy = None
        
    def set_strategy(self, strategy):
        """设置策略"""
        self.strategy = strategy
        self.strategy.set_backtest_mode(True)
        self.strategy.set_order_manager(self.order_manager)
        self.strategy.set_account_manager(self.account_manager)
        
    def load_data(self, exchange: str, symbol: str, interval: str, 
                  trade_type: str = 'spot') -> pd.DataFrame:
        """
        加载回测数据
        
        Args:
            exchange: 交易所
            symbol: 交易对
            interval: 时间间隔
            trade_type: 交易类型
            
        Returns:
            数据DataFrame
        """
        try:
            start_date = TimeUtils.parse_datetime(self.start_date) if isinstance(self.start_date, str) else self.start_date
            end_date = TimeUtils.parse_datetime(self.end_date) if isinstance(self.end_date, str) else self.end_date
            
            data = self.data_manager.get_data(
                exchange=exchange,
                symbol=symbol,
                interval=interval,
                trade_type=trade_type,
                start_date=start_date,
                end_date=end_date
            )
            
            if data.empty:
                self.logger.warning(f"未找到数据: {exchange} {symbol} {interval}")
                return pd.DataFrame()
            
            # 确保数据格式正确
            required_columns = ['open', 'high', 'low', 'close', 'volume']
            for col in required_columns:
                if col not in data.columns:
                    self.logger.error(f"缺少必要列: {col}")
                    return pd.DataFrame()
            
            # 确保时间索引
            if 'datetime' in data.columns and not isinstance(data.index, pd.DatetimeIndex):
                data.set_index('datetime', inplace=True)
            
            data.sort_index(inplace=True)
            
            self.logger.info(f"加载数据成功: {len(data)} 条记录")
            return data
            
        except Exception as e:
            self.logger.error(f"加载数据失败: {e}")
            return pd.DataFrame()
    
    def run_backtest(self, data: pd.DataFrame, symbol: str) -> Dict[str, Any]:
        """
        运行回测
        
        Args:
            data: 回测数据
            symbol: 交易对
            
        Returns:
            回测结果
        """
        try:
            if data.empty:
                raise ValueError("回测数据为空")
            
            if not self.strategy:
                raise ValueError("未设置策略")
            
            self.logger.info(f"开始回测: {symbol}, 数据量: {len(data)}")
            
            # 初始化策略
            self.strategy.initialize()
            
            # 逐条处理数据
            total_bars = len(data)
            for i, (timestamp, row) in enumerate(data.iterrows()):
                self.current_time = timestamp
                self.current_data[symbol] = row.to_dict()
                
                # 更新账户状态
                self._update_account_status(symbol, row)
                
                # 处理订单
                self._process_orders(symbol, row)
                
                # 调用策略
                try:
                    self.strategy.on_tick(symbol, row.to_dict())
                except Exception as e:
                    self.logger.error(f"策略执行错误 {timestamp}: {e}")
                    continue
                
                # 记录状态
                self._record_status(timestamp, symbol, row)
                
                # 进度显示
                if i % max(1, total_bars // 100) == 0:
                    progress = (i + 1) / total_bars * 100
                    self.logger.info(f"回测进度: {progress:.1f}%")
            
            # 计算最终指标
            self._calculate_metrics()
            
            self.logger.info("回测完成")
            return self.results
            
        except Exception as e:
            self.logger.error(f"回测失败: {e}")
            return {}
    
    def _update_account_status(self, symbol: str, bar: pd.Series):
        """更新账户状态"""
        try:
            current_price = bar['close']
            
            # 更新持仓价值
            if symbol in self.positions:
                position = self.positions[symbol]
                position['current_price'] = current_price
                position['market_value'] = position['quantity'] * current_price
                position['unrealized_pnl'] = (current_price - position['avg_price']) * position['quantity']
            
            # 更新账户总价值
            total_value = self.account_manager.get_balance()
            for pos in self.positions.values():
                total_value += pos.get('market_value', 0)
            
            self.account_manager.total_value = total_value
            
        except Exception as e:
            self.logger.error(f"更新账户状态失败: {e}")
    
    def _process_orders(self, symbol: str, bar: pd.Series):
        """处理订单"""
        try:
            current_price = bar['close']
            high_price = bar['high']
            low_price = bar['low']
            
            # 获取待处理订单
            pending_orders = self.order_manager.get_pending_orders()
            
            for order in pending_orders:
                if order['symbol'] != symbol:
                    continue
                
                # 检查订单是否可以成交
                can_fill = False
                fill_price = current_price
                
                if order['type'] == 'market':
                    can_fill = True
                    # 考虑滑点
                    if order['side'] == 'buy':
                        fill_price = current_price * (1 + self.slippage_rate)
                    else:
                        fill_price = current_price * (1 - self.slippage_rate)
                
                elif order['type'] == 'limit':
                    if order['side'] == 'buy' and low_price <= order['price']:
                        can_fill = True
                        fill_price = min(order['price'], current_price)
                    elif order['side'] == 'sell' and high_price >= order['price']:
                        can_fill = True
                        fill_price = max(order['price'], current_price)
                
                elif order['type'] == 'stop':
                    if order['side'] == 'buy' and high_price >= order['price']:
                        can_fill = True
                        fill_price = max(order['price'], current_price)
                    elif order['side'] == 'sell' and low_price <= order['price']:
                        can_fill = True
                        fill_price = min(order['price'], current_price)
                
                # 执行成交
                if can_fill:
                    self._fill_order(order, fill_price)
            
        except Exception as e:
            self.logger.error(f"处理订单失败: {e}")
    
    def _fill_order(self, order: Dict[str, Any], fill_price: float):
        """执行订单成交"""
        try:
            symbol = order['symbol']
            side = order['side']
            quantity = order['quantity']
            
            # 计算手续费
            commission = fill_price * quantity * self.commission_rate
            
            # 更新账户余额
            if side == 'buy':
                cost = fill_price * quantity + commission
                if self.account_manager.get_balance() >= cost:
                    self.account_manager.update_balance(-cost)
                    
                    # 更新持仓
                    if symbol not in self.positions:
                        self.positions[symbol] = {
                            'quantity': 0,
                            'avg_price': 0,
                            'total_cost': 0
                        }
                    
                    pos = self.positions[symbol]
                    total_quantity = pos['quantity'] + quantity
                    total_cost = pos['total_cost'] + fill_price * quantity
                    
                    pos['quantity'] = total_quantity
                    pos['avg_price'] = total_cost / total_quantity if total_quantity > 0 else 0
                    pos['total_cost'] = total_cost
                    
                else:
                    self.logger.warning(f"余额不足，无法买入: {order}")
                    return
            
            else:  # sell
                if symbol in self.positions and self.positions[symbol]['quantity'] >= quantity:
                    revenue = fill_price * quantity - commission
                    self.account_manager.update_balance(revenue)
                    
                    # 更新持仓
                    pos = self.positions[symbol]
                    pos['quantity'] -= quantity
                    pos['total_cost'] -= pos['avg_price'] * quantity
                    
                    if pos['quantity'] <= 0:
                        del self.positions[symbol]
                    
                else:
                    self.logger.warning(f"持仓不足，无法卖出: {order}")
                    return
            
            # 记录交易
            trade = {
                'timestamp': self.current_time,
                'symbol': symbol,
                'side': side,
                'quantity': quantity,
                'price': fill_price,
                'commission': commission,
                'order_id': order['id']
            }
            self.results['trades'].append(trade)
            
            # 更新订单状态
            self.order_manager.fill_order(order['id'], fill_price, quantity)
            
            self.logger.debug(f"订单成交: {trade}")
            
        except Exception as e:
            self.logger.error(f"订单成交失败: {e}")
    
    def _record_status(self, timestamp: datetime, symbol: str, bar: pd.Series):
        """记录状态"""
        try:
            # 记录余额历史
            balance_record = {
                'timestamp': timestamp,
                'balance': self.account_manager.get_balance(),
                'total_value': self.account_manager.total_value,
                'positions': self.positions.copy()
            }
            self.results['balance_history'].append(balance_record)
            
            # 记录持仓
            for symbol, pos in self.positions.items():
                position_record = {
                    'timestamp': timestamp,
                    'symbol': symbol,
                    'quantity': pos['quantity'],
                    'avg_price': pos['avg_price'],
                    'current_price': pos.get('current_price', 0),
                    'market_value': pos.get('market_value', 0),
                    'unrealized_pnl': pos.get('unrealized_pnl', 0)
                }
                self.results['positions'].append(position_record)
            
        except Exception as e:
            self.logger.error(f"记录状态失败: {e}")
    
    def _calculate_metrics(self):
        """计算回测指标"""
        try:
            if not self.results['balance_history']:
                return
            
            # 转换为DataFrame便于计算
            balance_df = pd.DataFrame(self.results['balance_history'])
            balance_df.set_index('timestamp', inplace=True)
            
            trades_df = pd.DataFrame(self.results['trades'])
            
            # 基本指标
            initial_value = self.initial_balance
            final_value = balance_df['total_value'].iloc[-1]
            total_return = (final_value - initial_value) / initial_value
            
            # 计算日收益率
            balance_df['daily_return'] = balance_df['total_value'].pct_change()
            
            # 年化收益率
            days = (balance_df.index[-1] - balance_df.index[0]).days
            annual_return = (1 + total_return) ** (365 / days) - 1 if days > 0 else 0
            
            # 最大回撤
            balance_df['cummax'] = balance_df['total_value'].cummax()
            balance_df['drawdown'] = (balance_df['total_value'] - balance_df['cummax']) / balance_df['cummax']
            max_drawdown = balance_df['drawdown'].min()
            
            # 夏普比率
            risk_free_rate = 0.02  # 假设无风险利率2%
            excess_return = annual_return - risk_free_rate
            volatility = balance_df['daily_return'].std() * np.sqrt(252)
            sharpe_ratio = excess_return / volatility if volatility > 0 else 0
            
            # 交易统计
            total_trades = len(trades_df)
            if total_trades > 0:
                winning_trades = len(trades_df[trades_df['side'] == 'sell'])  # 简化计算
                win_rate = winning_trades / total_trades if total_trades > 0 else 0
            else:
                win_rate = 0
            
            # 保存指标
            self.results['metrics'] = {
                'initial_value': initial_value,
                'final_value': final_value,
                'total_return': total_return,
                'annual_return': annual_return,
                'max_drawdown': max_drawdown,
                'sharpe_ratio': sharpe_ratio,
                'volatility': volatility,
                'total_trades': total_trades,
                'win_rate': win_rate,
                'days': days
            }
            
            self.logger.info(f"回测指标计算完成: {self.results['metrics']}")
            
        except Exception as e:
            self.logger.error(f"计算回测指标失败: {e}")
    
    def generate_report(self, output_path: str = None) -> str:
        """
        生成回测报告
        
        Args:
            output_path: 输出路径
            
        Returns:
            报告内容
        """
        try:
            metrics = self.results['metrics']
            
            report = f"""
回测报告
========

基本信息
--------
初始资金: ${metrics.get('initial_value', 0):,.2f}
最终资金: ${metrics.get('final_value', 0):,.2f}
回测天数: {metrics.get('days', 0)} 天

收益指标
--------
总收益率: {metrics.get('total_return', 0):.2%}
年化收益率: {metrics.get('annual_return', 0):.2%}
最大回撤: {metrics.get('max_drawdown', 0):.2%}

风险指标
--------
夏普比率: {metrics.get('sharpe_ratio', 0):.2f}
波动率: {metrics.get('volatility', 0):.2%}

交易统计
--------
总交易次数: {metrics.get('total_trades', 0)}
胜率: {metrics.get('win_rate', 0):.2%}
"""
            
            if output_path:
                Path(output_path).parent.mkdir(parents=True, exist_ok=True)
                with open(output_path, 'w', encoding='utf-8') as f:
                    f.write(report)
                self.logger.info(f"报告已保存: {output_path}")
            
            return report
            
        except Exception as e:
            self.logger.error(f"生成报告失败: {e}")
            return ""
    
    def plot_results(self, output_path: str = None):
        """
        绘制回测结果图表
        
        Args:
            output_path: 输出路径
        """
        try:
            if not self.results['balance_history']:
                self.logger.warning("无回测数据，无法绘图")
                return
            
            # 准备数据
            balance_df = pd.DataFrame(self.results['balance_history'])
            balance_df.set_index('timestamp', inplace=True)
            
            # 创建图表
            fig, axes = plt.subplots(2, 2, figsize=(15, 10))
            fig.suptitle('回测结果分析', fontsize=16)
            
            # 资金曲线
            axes[0, 0].plot(balance_df.index, balance_df['total_value'])
            axes[0, 0].set_title('资金曲线')
            axes[0, 0].set_ylabel('总资产')
            axes[0, 0].grid(True)
            
            # 回撤曲线
            balance_df['cummax'] = balance_df['total_value'].cummax()
            balance_df['drawdown'] = (balance_df['total_value'] - balance_df['cummax']) / balance_df['cummax']
            axes[0, 1].fill_between(balance_df.index, balance_df['drawdown'], 0, alpha=0.3, color='red')
            axes[0, 1].set_title('回撤曲线')
            axes[0, 1].set_ylabel('回撤比例')
            axes[0, 1].grid(True)
            
            # 日收益率分布
            balance_df['daily_return'] = balance_df['total_value'].pct_change()
            axes[1, 0].hist(balance_df['daily_return'].dropna(), bins=50, alpha=0.7)
            axes[1, 0].set_title('日收益率分布')
            axes[1, 0].set_xlabel('日收益率')
            axes[1, 0].set_ylabel('频次')
            axes[1, 0].grid(True)
            
            # 月度收益
            monthly_returns = balance_df['total_value'].resample('M').last().pct_change()
            axes[1, 1].bar(range(len(monthly_returns)), monthly_returns.values)
            axes[1, 1].set_title('月度收益率')
            axes[1, 1].set_xlabel('月份')
            axes[1, 1].set_ylabel('月收益率')
            axes[1, 1].grid(True)
            
            plt.tight_layout()
            
            if output_path:
                Path(output_path).parent.mkdir(parents=True, exist_ok=True)
                plt.savefig(output_path, dpi=300, bbox_inches='tight')
                self.logger.info(f"图表已保存: {output_path}")
            
            plt.show()
            
        except Exception as e:
            self.logger.error(f"绘制图表失败: {e}")
    
    def save_results(self, output_path: str):
        """
        保存回测结果
        
        Args:
            output_path: 输出路径
        """
        try:
            Path(output_path).parent.mkdir(parents=True, exist_ok=True)
            
            # 转换datetime对象为字符串
            results_copy = self.results.copy()
            
            # 处理balance_history
            for record in results_copy['balance_history']:
                record['timestamp'] = record['timestamp'].isoformat()
            
            # 处理trades
            for trade in results_copy['trades']:
                trade['timestamp'] = trade['timestamp'].isoformat()
            
            # 处理positions
            for position in results_copy['positions']:
                position['timestamp'] = position['timestamp'].isoformat()
            
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(results_copy, f, indent=2, ensure_ascii=False)
            
            self.logger.info(f"回测结果已保存: {output_path}")
            
        except Exception as e:
            self.logger.error(f"保存结果失败: {e}")